package com.fengwk.cv4j.view.ind;

import org.nd4j.linalg.api.ndarray.INDArray;

import com.fengwk.cv4j.compute.Compute;
import com.fengwk.cv4j.factory.ViewFactory;
import com.fengwk.cv4j.utils.Nd4jUtils;
import com.fengwk.cv4j.view.View;
import com.fengwk.cv4j.view.common.MaxView;
import com.fengwk.support.exception.BaseException;
import com.fengwk.support.function.Return;

/**
 * 取矩阵某维度的最大值
 * 
 * @author fengwk
 *
 */
public class INDMaxView extends MaxView<INDArrayWrap> {

    private static final long serialVersionUID = -8788017840837179680L;

    public INDMaxView(Compute<INDArrayWrap> compute, ViewFactory<INDArrayWrap> factory, @SuppressWarnings("unchecked") View<INDArrayWrap>...vs) {
        super(compute, factory, vs);
    }

    @Override
    protected View<INDArrayWrap> gradient0(int idx) {
    	View<INDArrayWrap> d = gradient0(idx, len(), ts(), f(), (a, ds) -> Nd4jUtils.max0(a, ds));
    	if (d == null)
    		return super.gradient0(idx);
    	return d;
    }
    
    static View<INDArrayWrap> gradient0(int idx, int len, INDArrayWrap[] ts, ViewFactory<INDArrayWrap> f, Return.Params2<INDArray, INDArray, int[]> computeOp) {
    	INDArrayWrap wrap0 = ts[0];
        INDArray arr = wrap0.arr();
        int[] dimension = new int[len - 1];
        for (int i = 1; i < len; i ++) {
            Number num = ts[i].num();
            if (num == null)
                throw new BaseException("num == null");
            dimension[i - 1] = num.intValue();
        }
        if (arr != null) {
            if (idx == 0) {
            	INDArray max = computeOp.apply(arr, dimension);
            	if (max == null)
            		throw new BaseException("max == null");
            	INDArray molecular = arr.dup();
            	max = max.broadcast(molecular.shape());
                // 将原始矩阵化为0和1的矩阵为分子
                for (int i = 0; i < molecular.length(); i ++) {
                	double b = molecular.getDouble(i), m = max.getDouble(i);
                	if (b != m)
                		molecular.putScalar(i, 0);
                	else
                		molecular.putScalar(i, 1);
                }
                // 这个0和1的矩阵按max维度求和并广播后将成为父母
                INDArray denominator = Nd4jUtils.sum0(molecular, dimension).broadcast(molecular.shape());
                // 修正父母0的除法
                for (int i = 0; i < molecular.length(); i ++) {
                	double m = molecular.getDouble(i), d = denominator.getDouble(i);
                	if (m == 0)
                		molecular.putScalar(i, 0);
                	else
                		molecular.putScalar(i, m / d);
                }
                return f.of(new INDArrayWrap(molecular));
            }
            if (idx < len)
                return f.of(0);
            throw new BaseException("idx >= len()"); 
        }
		return null;
    }

}
